#include "utils.h"
#include <openssl/hmac.h>
#include <openssl/evp.h>
#include <openssl/bio.h>
#include <openssl/buffer.h>
#include <string.h>
#include <stdlib.h>
#include <stdio.h>

const char* secret = "my_secure_secret_key";

// Base64 编码函数
char* base64_encode(const unsigned char* input, int length) {
    BIO *bio, *b64;
    BUF_MEM *bufferPtr;

    b64 = BIO_new(BIO_f_base64());
    bio = BIO_new(BIO_s_mem());
    bio = BIO_push(b64, bio);

    // 不添加换行符
    BIO_set_flags(bio, BIO_FLAGS_BASE64_NO_NL);
    BIO_write(bio, input, length);
    BIO_flush(bio);
    BIO_get_mem_ptr(bio, &bufferPtr);
    
    // 分配内存并复制数据
    char* output = (char*)malloc(bufferPtr->length + 1);
    memcpy(output, bufferPtr->data, bufferPtr->length);
    output[bufferPtr->length] = '\0';
    
    BIO_free_all(bio);
    return output;
}

// Base64 解码函数
unsigned char* base64_decode(const char* input, int* out_length) {
    BIO *bio, *b64;
    int length = strlen(input);
    
    // 分配足够大的输出缓冲区
    unsigned char* buffer = (unsigned char*)malloc(length);
    memset(buffer, 0, length);
    
    b64 = BIO_new(BIO_f_base64());
    bio = BIO_new_mem_buf(input, length);
    bio = BIO_push(b64, bio);
    
    // 不期待换行符
    BIO_set_flags(bio, BIO_FLAGS_BASE64_NO_NL);
    *out_length = BIO_read(bio, buffer, length);
    
    BIO_free_all(bio);
    return buffer;
}

// URL 安全的 Base64 编码
char* base64_url_encode(const unsigned char* input, int length) {
    char* base64 = base64_encode(input, length);
    // 替换Base64中的特殊字符
    for (int i = 0; base64[i] != '\0'; i++) {
        if (base64[i] == '+') base64[i] = '-';
        if (base64[i] == '/') base64[i] = '_';
        if (base64[i] == '=') base64[i] = '\0'; // 移除填充
    }
    return base64;
}

// URL 安全的 Base64 解码
unsigned char* base64_url_decode(const char* input, int* out_length) {
    int length = strlen(input);
    char* padded = (char*)malloc(length + 4); // 为填充预留空间
    strcpy(padded, input);
    
    // 恢复Base64字符
    for (int i = 0; padded[i] != '\0'; i++) {
        if (padded[i] == '-') padded[i] = '+';
        if (padded[i] == '_') padded[i] = '/';
    }
    
    // 添加填充
    int pad_len = (4 - (length % 4)) % 4;
    for (int i = 0; i < pad_len; i++) {
        strcat(padded, "=");
    }
    
    unsigned char* result = base64_decode(padded, out_length);
    free(padded);
    return result;
}

// HMAC-SHA256 签名
char* hmac_sha256_sign(const char* key, const char* data) {
    unsigned char* digest = HMAC(EVP_sha256(), 
                                key, strlen(key),
                                (const unsigned char*)data, strlen(data),
                                NULL, NULL);
    return base64_url_encode(digest, 32); // SHA256摘要长度是32字节
}

// 生成JWT Token
char* generate_jwt_token(const char* secret, int user_id, const char* username, 
                         const char* avatar_initials, const char* avatar_gradient) {
    // 1. 构建Header
    const char* header = "{\"alg\":\"HS256\",\"typ\":\"JWT\"}";
    char* header_b64 = base64_url_encode((const unsigned char*)header, strlen(header));
    
    // 2. 构建Payload
    time_t now = time(NULL);
    struct json_object* payload = json_object_new_object();
    json_object_object_add(payload, "user_id", json_object_new_int(user_id));
    json_object_object_add(payload, "username", json_object_new_string(username));
    json_object_object_add(payload, "avatar_initials", json_object_new_string(avatar_initials));
    json_object_object_add(payload, "avatar_gradient", json_object_new_string(avatar_gradient));
    json_object_object_add(payload, "iat", json_object_new_int64(now));
    json_object_object_add(payload, "exp", json_object_new_int64(now + 1800)); // 30分钟有效期
    
    const char* payload_str = json_object_to_json_string(payload);
    char* payload_b64 = base64_url_encode((const unsigned char*)payload_str, strlen(payload_str));
    
    // 3. 生成签名
    char* data = (char*)malloc(strlen(header_b64) + strlen(payload_b64) + 2);
    sprintf(data, "%s.%s", header_b64, payload_b64);
    
    char* signature = hmac_sha256_sign(secret, data);
    
    // 4. 组合Token
    char* token = (char*)malloc(strlen(header_b64) + strlen(payload_b64) + strlen(signature) + 3);
    sprintf(token, "%s.%s.%s", header_b64, payload_b64, signature);
    
    // 清理
    free(header_b64);
    json_object_put(payload);
    free(payload_b64);
    free(data);
    free(signature);
    
    return token;
}

// 验证JWT Token
int verify_jwt_token(const char* token, const char* secret) {
    // 1. 分割Token
    char* token_copy = strdup(token);
    char* parts[3];
    char* part = strtok(token_copy, ".");
    int i = 0;
    
    while (part != NULL && i < 3) {
        parts[i++] = part;
        part = strtok(NULL, ".");
    }
    
    if (i != 3) {
        free(token_copy);
        return 0; // 无效的Token格式
    }
    
    // 2. 重新生成签名
    char* data = (char*)malloc(strlen(parts[0]) + strlen(parts[1]) + 2);
    sprintf(data, "%s.%s", parts[0], parts[1]);
    
    char* real_signature = hmac_sha256_sign(secret, data);
    
    // 3. 比较签名
    int valid = (strcmp(parts[2], real_signature) == 0);
    
    // 4. 检查过期时间
    if (valid) {
        int out_len;
        unsigned char* payload_bin = base64_url_decode(parts[1], &out_len);
        payload_bin[out_len] = '\0'; // 确保字符串终止
        
        // 解析JSON
        struct json_object* payload = json_tokener_parse((char*)payload_bin);
        if (payload) {
            struct json_object* exp_obj;
            if (json_object_object_get_ex(payload, "exp", &exp_obj)) {
                time_t exp = json_object_get_int64(exp_obj);
                if (time(NULL) > exp) {
                    valid = 0; // Token已过期
                }
            } else {
                valid = 0; // 没有exp字段
            }
            json_object_put(payload);
        } else {
            valid = 0; // 无效的Payload
        }
        free(payload_bin);
    }
    
    // 清理
    free(token_copy);
    free(data);
    free(real_signature);
    
    return valid;
}

// 从JWT Token中解析Payload
struct json_object* parse_jwt_payload(const char* token) {
    // 1. 分割Token
    char* token_copy = strdup(token);
    char* parts[3];
    char* part = strtok(token_copy, ".");
    int i = 0;
    
    while (part != NULL && i < 3) {
        parts[i++] = part;
        part = strtok(NULL, ".");
    }
    
    if (i < 2) {
        free(token_copy);
        return NULL;
    }
    
    // 2. 解码Payload
    int out_len;
    unsigned char* payload_bin = base64_url_decode(parts[1], &out_len);
    payload_bin[out_len] = '\0'; // 确保字符串终止
    
    // 3. 解析为JSON
    struct json_object* payload = json_tokener_parse((char*)payload_bin);
    
    // 清理
    free(token_copy);
    free(payload_bin);
    
    return payload;
}

// 从Payload中获取用户ID
int get_user_id_from_payload(struct json_object* payload) {
    struct json_object* user_id_obj;
    if (json_object_object_get_ex(payload, "user_id", &user_id_obj)) {
        return json_object_get_int(user_id_obj);
    }
    return -1;
}

// 从Payload中获取用户名
const char* get_username_from_payload(struct json_object* payload) {
    struct json_object* username_obj;
    if (json_object_object_get_ex(payload, "username", &username_obj)) {
        return json_object_get_string(username_obj);
    }
    return NULL;
}

// 从Payload中获取头像信息
void get_avatar_info_from_payload(struct json_object* payload, const char** initials, const char** gradient) {
    struct json_object* initials_obj;
    struct json_object* gradient_obj;
    
    if (json_object_object_get_ex(payload, "avatar_initials", &initials_obj)) {
        *initials = json_object_get_string(initials_obj);
    }
    
    if (json_object_object_get_ex(payload, "avatar_gradient", &gradient_obj)) {
        *gradient = json_object_get_string(gradient_obj);
    }
}

// 检查Token是否即将过期
int is_token_expiring_soon(const char* token, int threshold_seconds) {
    struct json_object* payload = parse_jwt_payload(token);
    if (!payload) return 1;
    
    struct json_object* exp_obj;
    if (!json_object_object_get_ex(payload, "exp", &exp_obj)) {
        json_object_put(payload);
        return 1;
    }
    
    time_t exp = json_object_get_int64(exp_obj);
    time_t now = time(NULL);
    
    json_object_put(payload);
    return (exp - now) < threshold_seconds;
}

void generate_secret_key(char *key, int length) {
    const char charset[] = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789";
    if (length) {
        --length;
        for (int n = 0; n < length; n++) {
            int index = rand() % (int)(sizeof charset - 1);
            key[n] = charset[index];
        }
        key[length] = '\0';
    }
}